#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
交互式频散曲线拾取程序 PyQt6 版本
支持从FV图中人工拾取频散曲线，并自动调整到局部极大值点
"""

import sys
import os
import numpy as np
import h5py
from datetime import datetime

from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QFileDialog, QMessageBox, QSizePolicy,
                             QButtonGroup, QPushButton, QGroupBox, QGridLayout,
                             QLabel, QDialog, QFormLayout, QLineEdit, QDialogButtonBox)
from PyQt6.QtCore import Qt, pyqtSignal
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure

def norm(data):
    """归一化数据"""
    data = np.array(data)
    if np.max(np.abs(data)) == 0:
        return data
    return data / np.max(np.abs(data))


def do_fv(data, vs, x=None, sx=0, dt=1e-2, dx=50, NORM=False, PRINT=False, fs=-1, fe=1e10):
    """计算FV数据"""
    I = np.sqrt(np.array(-1, dtype='complex'))
    PI = np.pi

    nx, nt = data.shape
    if x is None:
        x = np.arange(nx) * dx
    t = np.arange(nt) * dt

    XF = np.zeros([nx, nt], dtype='complex')
    for ix in range(nx):
        XF[ix, :] = np.fft.fft(data[ix, :])
        if PRINT:
            print('\r FX:{}|{}'.format(ix, nx), end='', flush=True)
    FX = XF.T
    f = np.arange(nt) * 1 / t.max()
    if PRINT:
        print('')
    nv = len(vs)
    FV = np.zeros([nt, nv], dtype='complex')
    idx_valid = []
    for j in range(nt // 2):
        if f[j] < fs or f[j] > fe:
            continue
        idx_valid.append(j)
        for i, v in enumerate(vs):
            # FV[j,i] = np.sum(np.exp(I*2*PI*np.abs(x-sx)/v*f[j])*FX[j,:])
            FV[j, i] = np.sum(np.exp(I * 2 * PI * (x - sx) / v * f[j]) * FX[j, :])
        if NORM:
            FV[j, :] = FV[j, :] / np.abs(FV[j, :]).max()
        if PRINT:
            print('\r FV:{}|{}'.format(j, nt), end='')

    return FV[idx_valid, :], f[idx_valid], vs


class FVConfigDialog(QDialog):
    """FV参数配置对话框"""
    def __init__(self, parent=None, vs_min=100, vs_max=1400, f_min=1.0, f_max=5.0):
        super().__init__(parent)
        self.setWindowTitle("FV Configuration")
        
        layout = QFormLayout(self)
        
        self.vs_min_edit = QLineEdit(str(vs_min))
        self.vs_max_edit = QLineEdit(str(vs_max))
        self.f_min_edit = QLineEdit(str(f_min))
        self.f_max_edit = QLineEdit(str(f_max))
        
        layout.addRow("Vs Min (m/s):", self.vs_min_edit)
        layout.addRow("Vs Max (m/s):", self.vs_max_edit)
        layout.addRow("F Min (Hz):", self.f_min_edit)
        layout.addRow("F Max (Hz):", self.f_max_edit)
        
        button_box = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel
        )
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        layout.addRow(button_box)
        
    def get_values(self):
        """获取配置值"""
        return {
            'vs_min': float(self.vs_min_edit.text()),
            'vs_max': float(self.vs_max_edit.text()),
            'f_min': float(self.f_min_edit.text()),
            'f_max': float(self.f_max_edit.text())
        }


class DispersionPickerWindow(QMainWindow):
    def __init__(self, h5_file=None):
        super().__init__()
        self.setWindowTitle('Dispersion Curve Picker')
        # 调整窗口大小，使其更扁平
        self.setGeometry(100, 100, 1800, 800)
        
        # 配置参数
        self.CONFIG = {
            'DPI': 100,
            'FIGSIZE': (16, 7),  # 调整图形大小比例
            'SEARCH_RANGE': 50,  # 速度搜索范围（米/秒）
            'DOUBLE_CLICK_SEARCH_RANGE': 20,  # 双击时的速度搜索范围（米/秒）
            'FV_NORM': True,
            'DEFAULT_SCALE': 60
        }
        
        # 初始化默认参数
        self.vs = np.linspace(100, 1400, 800)  # 默认速度范围
        self.fs = 1.0  # 默认频率下限
        self.fe = 5.0  # 默认频率上限
        self.SCALE = 60  # 波形显示的SCALE初始值
        
        # 初始化变量
        self.h5_file = h5_file
        self.have_fv_data = False
        self.current_curve = None  # 当前拾取的曲线
        self.current_curve_index = 0  # 当前曲线索引（支持多条曲线）
        self.curves = {}  # 存储所有拾取的曲线
        self.selected_point = None  # 当前选中的点
        self.selected_point_marker = None  # 用于标记选中点的图形对象
        self.curve_lines = {}  # 存储曲线线条对象
        self.point_annotations = {}  # 存储点标注对象
        self.mode = 'edit'  # 统一为edit模式
        
        # 创建主部件和布局
        self.main_widget = QWidget(self)
        self.setCentralWidget(self.main_widget)
        self.main_layout = QHBoxLayout(self.main_widget)
        
        # 创建菜单栏
        self.create_menu_bar()
        
        # 创建控制面板
        self.create_control_panel()
        
        # 创建图形区域
        self.create_figure_area()
        
        # 创建状态栏
        self.statusBar().showMessage("Ready")
        self.status_label = QLabel("Current Curve: 0")
        self.h5_file_label = QLabel("H5 File: " + (os.path.basename(self.h5_file) if self.h5_file else "None"))
        self.statusBar().addPermanentWidget(self.status_label)
        # 添加弹簧以将H5文件标签推到中间
        self.statusBar().addPermanentWidget(QWidget())
        self.statusBar().addPermanentWidget(self.h5_file_label)
        
        # 如果提供了h5文件，则加载数据
        if self.h5_file and os.path.exists(self.h5_file):
            self.load_h5_data()
            self.plot_fv_and_waveforms()
            # 尝试加载对应的.disp文件
            self.try_load_disp_file()
            # 如果没有加载到曲线，则初始化参考曲线
            if not self.curves:
                self.initialize_reference_curve()
        
    def create_figure_area(self):
        """创建图形显示区域"""
        # 创建图形对象
        self.figure = Figure(figsize=self.CONFIG['FIGSIZE'], dpi=self.CONFIG['DPI'])
        self.canvas = FigureCanvas(self.figure)
        self.canvas.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        
        # 创建导航工具栏
        self.toolbar = NavigationToolbar(self.canvas, self)
        
        # 创建垂直布局用于图形和工具栏
        figure_layout = QVBoxLayout()
        figure_layout.addWidget(self.toolbar)
        figure_layout.addWidget(self.canvas)
        
        # 创建包含图形的部件
        figure_widget = QWidget()
        figure_widget.setLayout(figure_layout)
        
        # 添加到主布局
        self.main_layout.addWidget(figure_widget)
        
    def create_menu_bar(self):
        """创建菜单栏"""
        menubar = self.menuBar()
        
        # 文件菜单
        file_menu = menubar.addMenu('File')
        
        open_action = file_menu.addAction('Open H5 File')
        open_action.triggered.connect(self.open_h5_file)
        
    def create_control_panel(self):
        """创建控制面板"""
        # 创建控制面板部件
        self.control_panel = QGroupBox("Controls")
        self.control_panel.setMaximumWidth(200)
        control_layout = QVBoxLayout()
        
        # 模式按钮组
        mode_group = QGroupBox("Mode")
        mode_layout = QVBoxLayout()
        
        self.btn_group_mode = QButtonGroup()
        
        self.btn_edit_mode = QPushButton("Edit Mode")
        self.btn_edit_mode.setCheckable(True)
        self.btn_edit_mode.setChecked(True)
        self.btn_group_mode.addButton(self.btn_edit_mode)
        
        mode_layout.addWidget(self.btn_edit_mode)
        mode_group.setLayout(mode_layout)
        control_layout.addWidget(mode_group)
        
        # 曲线操作按钮组
        curve_group = QGroupBox("Curve Operations")
        curve_layout = QVBoxLayout()
        
        self.btn_new_curve = QPushButton("New Curve")
        self.btn_delete_curve = QPushButton("Delete Curve")
        self.btn_redo_fv = QPushButton("Redo FV")
        
        self.btn_new_curve.clicked.connect(self.new_curve)
        self.btn_delete_curve.clicked.connect(self.delete_current_curve)
        self.btn_redo_fv.clicked.connect(self.show_fv_config_dialog)
        
        curve_layout.addWidget(self.btn_new_curve)
        curve_layout.addWidget(self.btn_delete_curve)
        curve_layout.addWidget(self.btn_redo_fv)
        curve_group.setLayout(curve_layout)
        control_layout.addWidget(curve_group)
        
        # 文件操作按钮组
        file_group = QGroupBox("File Operations")
        file_layout = QVBoxLayout()
        
        self.btn_open_file = QPushButton("Open H5 File")
        self.btn_save_curves = QPushButton("Save Curves")
        self.btn_load_curves = QPushButton("Load Curves")
        self.btn_save_figure = QPushButton("Save Figure (.disp.png)")
        
        self.btn_open_file.clicked.connect(self.open_h5_file)
        self.btn_save_curves.clicked.connect(self.save_curves)
        self.btn_load_curves.clicked.connect(self.load_curves_from_file)
        self.btn_save_figure.clicked.connect(self.save_figure)
        
        file_layout.addWidget(self.btn_open_file)
        file_layout.addWidget(self.btn_save_curves)
        file_layout.addWidget(self.btn_load_curves)
        file_layout.addWidget(self.btn_save_figure)
        file_group.setLayout(file_layout)
        control_layout.addWidget(file_group)
        
        # 曲线切换按钮组
        switch_group = QGroupBox("Switch Curve")
        switch_layout = QGridLayout()
        
        # 定义曲线按钮的标签
        self.curve_button_labels = ['M0', 'M1', 'M2', 'M3', 'M4', 'M5']
        self.curve_buttons = []
        for i in range(6):  # 支持最多6条曲线
            btn = QPushButton(self.curve_button_labels[i])
            btn.clicked.connect(lambda checked, idx=i: self.switch_to_curve(idx))
            btn.setEnabled(False)
            # 设置按钮样式
            btn.setStyleSheet("color: black;")
            switch_layout.addWidget(btn, i // 3, i % 3)
            self.curve_buttons.append(btn)
            
        switch_group.setLayout(switch_layout)
        control_layout.addWidget(switch_group)
        
        # 添加弹簧以将控件推到顶部
        control_layout.addStretch()
        
        self.control_panel.setLayout(control_layout)
        self.main_layout.addWidget(self.control_panel)
        
    def load_h5_data(self):
        """从HDF5文件加载数据"""
        try:
            with h5py.File(self.h5_file, 'r') as f:
                # 尝试读取FV相关数据
                if 'FV' in f and 'vs' in f and 'f' in f:
                    self.FV = np.abs(f['FV'][()])  # 取模值
                    self.vs = f['vs'][()]
                    self.f = f['f'][()]
                    self.have_fv_data = True
                else:
                    self.have_fv_data = False
                    
                # 读取波形数据（必须存在）
                self.traces = f['traces'][()]
                self.x = f['x'][()]
                self.t = f['t'][()]
                
            # 如果没有FV数据，则计算默认的FV矩阵
            if not self.have_fv_data:
                print("FV data not found, will compute FV matrix with default parameters")
                # 这里暂时不实现FV计算，只处理已有的数据
                pass
            
            # 归一化FV图
            if hasattr(self, 'FV'):
                self.FV = self.FV / np.max(self.FV)
            
            # 保存原始数据范围
            self.original_vs = self.vs.copy()
            self.original_f = self.f.copy() if hasattr(self, 'f') else None
            if hasattr(self, 'FV'):
                self.original_FV = self.FV.copy()
            
            # 初始化变量
            self.nv = len(self.vs)
            self.nt = len(self.f) if hasattr(self, 'f') else 0
            self.dt = self.t[1] - self.t[0] if len(self.t) > 1 else 1e-2  # 时间采样间隔
            
            # 更新状态栏中的H5文件名
            self.h5_file_label.setText("H5 File: " + os.path.basename(self.h5_file))
            
        except Exception as e:
            print(f"Error loading H5 file: {e}")
            
    def plot_fv_and_waveforms(self):
        """绘制FV图和波形图"""
        # 清除之前的图形
        self.figure.clear()
        
        # 创建子图 - FV图在左侧，波形图在右侧
        self.ax_fv = self.figure.add_subplot(1, 2, 1)
        self.ax_waveform = self.figure.add_subplot(1, 2, 2)
        
        # 显示FV图 (f为横轴，v为纵轴，速度轴上大下小)
        if hasattr(self, 'FV') and hasattr(self, 'f') and hasattr(self, 'vs'):
            # 使用pcolormesh替代imshow
            F, V = np.meshgrid(self.f, self.vs, indexing='ij')
            self.im = self.ax_fv.pcolormesh(F, V, self.FV, shading='nearest', cmap='nipy_spectral_r')
            self.ax_fv.set_xlabel('Frequency (Hz)')
            self.ax_fv.set_ylabel('Phase Velocity (m/s)')
            self.ax_fv.set_title('F-V Spectrum')  # 标题保持简洁
            self.ax_fv.set_ylim(self.vs.min(), self.vs.max())
            self.ax_fv.set_xlim(self.f.min(), self.f.max())
            # 连接滚轮事件
            self.canvas.mpl_connect('scroll_event', self.on_scroll)
        
        # 显示波形数据
        if hasattr(self, 'traces') and len(self.traces) > 0:
            self.plot_waveforms()
        self.figure.tight_layout()
        # 连接鼠标事件
        self.canvas.mpl_connect('button_press_event', self.on_click)
        
        # 刷新图形
        self.canvas.draw()
        
    def plot_waveforms(self):
        """绘制所有波形数据"""
        # 清除之前的波形图
        self.ax_waveform.clear()
        
        # 使用滤波后的数据（如果存在），否则使用原始数据
        traces_to_plot = getattr(self, 'filtered_traces', self.traces)
        
        nx = len(traces_to_plot)
        if nx == 0:
            return
            
        for i in range(nx):
            trace_i = traces_to_plot[i, :].copy()
            trace_i = norm(trace_i)
            trace_i = trace_i * self.SCALE
            self.ax_waveform.plot(self.t, trace_i + self.x[i], lw=0.5, color='k')
            
        self.ax_waveform.set_xlabel('Time (s)')
        self.ax_waveform.set_ylabel('Distance (m)')
        self.ax_waveform.set_title('Waveforms')
        
        # 添加频率范围信息
        self.ax_waveform.text(0.02, 0.98, f'Freq Range: {self.fs:.1f}-{self.fe:.1f} Hz', 
                             transform=self.ax_waveform.transAxes, 
                             verticalalignment='top',
                             bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

    def initialize_reference_curve(self):
        """基于每个频率点上的最大能量初始化参考曲线"""
        # 如果没有FV数据则返回
        if not (hasattr(self, 'FV') and hasattr(self, 'f') and hasattr(self, 'vs')):
            return
            
        if len(self.FV) == 0 or len(self.f) == 0 or len(self.vs) == 0:
            return
            
        curve_points = []
        
        # 对每个频率找到最大能量对应的速度
        for i in range(len(self.f)):
            # 找到该频率下能量最大的速度点
            max_idx = np.argmax(self.FV[i, :])
            curve_points.append((self.f[i], self.vs[max_idx]))
        
        # 保存为第一条曲线
        self.curves[0] = np.array(curve_points)
        self.current_curve_index = 0
        self.current_curve = self.curves[0]
        self.draw_curve()
        
        # 更新曲线按钮状态
        self.update_curve_buttons()
        self.update_status()
        
    def update_curve_buttons(self):
        """更新曲线按钮状态"""
        # 启用已有的曲线按钮
        for i, btn in enumerate(self.curve_buttons):
            btn.setEnabled(i in self.curves or i == 0)  # 至少启用第一个
            # 根据是否有曲线数据更新按钮样式
            if i in self.curves:
                btn.setStyleSheet("font-weight: normal; color: black;")
            else:
                btn.setStyleSheet("color: gray;")  # 灰色表示无数据
                
        # 高亮当前活动曲线按钮（使用红色）
        if self.current_curve_index < len(self.curve_buttons):
            self.curve_buttons[self.current_curve_index].setStyleSheet("font-weight: bold; color: red;")
            
    def update_status(self):
        """更新状态显示"""
        self.status_label.setText(f"Current Curve: {self.current_curve_index}")
        
    def find_local_maximum(self, freq_val, initial_vs, is_double_click=False):
        """
        在给定频率下，找到initial_vs附近的局部能量极大值点
        支持对已有曲线点进行微调
        """
        # 确保我们有必要的数据
        if not (hasattr(self, 'f') and hasattr(self, 'vs') and hasattr(self, 'FV')):
            return initial_vs
            
        # 找到最接近的频率索引
        freq_idx = np.argmin(np.abs(self.f - freq_val))
        
        # 找到最接近的vs索引
        vs_idx = np.argmin(np.abs(self.vs - initial_vs))
        
        # 在该频率下寻找局部极大值
        fv_row = self.FV[freq_idx, :]
        
        # 根据是否是双击选择不同的搜索范围
        if is_double_click:
            search_window_m_s = self.CONFIG['DOUBLE_CLICK_SEARCH_RANGE']  # 双击时的搜索范围
        else:
            search_window_m_s = self.CONFIG['SEARCH_RANGE']  # 普通点击的搜索范围
            
        # 计算速度分辨率
        dv = (self.vs[-1] - self.vs[0]) / len(self.vs) if len(self.vs) > 1 else 1
        # 计算搜索窗口点数
        search_window = int(search_window_m_s / dv) if dv > 0 else 100
        search_window = max(1, search_window)  # 至少为1
        
        # 确保搜索窗口在有效范围内
        start_idx = max(0, vs_idx - search_window)
        end_idx = min(len(self.vs), vs_idx + search_window)
        
        # 在指定窗口内找到局部最大值
        local_max_idx = start_idx + np.argmax(fv_row[start_idx:end_idx])
        
        # 返回找到的最大值对应的速度
        return self.vs[local_max_idx]
        
    def draw_curve(self):
        """绘制当前曲线"""
        # 定义颜色列表，为不同曲线使用不同颜色
        colors = ['white', 'red', 'yellow', 'cyan', 'magenta', 'lime', 'orange', 'pink']
        
        # 清除之前的曲线和标注
        if self.current_curve_index in self.curve_lines:
            try:
                self.curve_lines[self.current_curve_index].remove()
            except (NotImplementedError, ValueError):
                # 如果无法直接移除，将其设为不可见
                self.curve_lines[self.current_curve_index].set_visible(False)
            # 同时清除标注
            for ann in self.point_annotations.get(self.current_curve_index, []):
                try:
                    ann.remove()
                except (NotImplementedError, ValueError):
                    # 如果无法直接移除，将其设为不可见
                    ann.set_visible(False)
            if self.current_curve_index in self.point_annotations:
                del self.point_annotations[self.current_curve_index]
            
        # 绘制当前曲线
        if self.current_curve is not None and len(self.current_curve) > 0:
            # 根据曲线索引选择颜色
            color_idx = self.current_curve_index % len(colors)
            curve_color = colors[color_idx]
            
            line, = self.ax_fv.plot(self.current_curve[:, 0], self.current_curve[:, 1], 
                                   'o-', color=curve_color, linewidth=2, markersize=4,
                                   label=f'Curve {self.current_curve_index}', picker=True)
            self.curve_lines[self.current_curve_index] = line
            
            # 添加标注
            self.point_annotations[self.current_curve_index] = []
            
            for i, (freq, vs) in enumerate(self.current_curve):
                ann = self.ax_fv.annotate(str(i), (freq, vs), 
                                         xytext=(5, 5), textcoords='offset points',
                                         color=curve_color, fontsize=8)
                self.point_annotations[self.current_curve_index].append(ann)
        
        self.ax_fv.legend()
        self.canvas.draw()
        
    def on_click(self, event):
        """处理鼠标点击事件"""
        # 只处理FV图上的点击
        if event.inaxes != self.ax_fv:
            return
            
        freq = event.xdata
        vs = event.ydata
        
        if freq is None or vs is None:
            return
            
        # 找到最近的频率索引（从self.f选择，不能自己添加更密集的点）
        freq_idx = np.argmin(np.abs(self.f - freq))
        nearest_freq = self.f[freq_idx]
        
        if event.button == 1 and event.dblclick:
            # 双击 - 缩小极大值范围
            adjusted_vs = self.find_local_maximum(nearest_freq, vs, is_double_click=True)
            self.handle_point_selection(nearest_freq, adjusted_vs)
        elif event.button == 1:  # 左键 - 添加或更新点
            adjusted_vs = self.find_local_maximum(nearest_freq, vs, is_double_click=False)
            self.handle_point_selection(nearest_freq, adjusted_vs)
        elif event.button == 3:  # 右键 - 删除离点击坐标最近的点
            self.remove_nearest_point(freq, vs)
                
    def handle_point_selection(self, freq, vs):
        """处理点选择（添加或更新）"""
        if self.current_curve is None:
            self.current_curve = np.empty((0, 2))
            
        # 检查是否已经有点在这个频率上
        freq_exists = False
        point_index = -1
        if self.current_curve is not None and len(self.current_curve) > 0:
            for i, (existing_freq, _) in enumerate(self.current_curve):
                if np.isclose(existing_freq, freq, atol=0.01):
                    # 更新现有频率点的速度
                    self.current_curve[i] = [freq, vs]
                    freq_exists = True
                    point_index = i
                    break
                
        if not freq_exists:
            # 添加新点
            new_point = np.array([[freq, vs]])
            self.current_curve = np.vstack([self.current_curve, new_point]) if len(self.current_curve) > 0 else new_point
            point_index = len(self.current_curve) - 1 if len(self.current_curve) > 0 else 0
            
        # 按频率排序
        if self.current_curve is not None and len(self.current_curve) > 0:
            sorted_indices = np.argsort(self.current_curve[:, 0])
            self.current_curve = self.current_curve[sorted_indices]
            
        # 更新存储
        self.curves[self.current_curve_index] = self.current_curve
        self.draw_curve()
        
        # 标记选中的点为红色高亮
        self.mark_selected_point(freq, vs)
        
    def mark_selected_point(self, freq, vs):
        """标记选中的点为红色"""
        # 清除之前的标记
        if self.selected_point_marker:
            try:
                self.selected_point_marker.remove()
            except (NotImplementedError, ValueError):
                self.selected_point_marker.set_visible(False)
            
        # 添加新的标记
        self.selected_point_marker, = self.ax_fv.plot(freq, vs, 'ro', markersize=8, markeredgewidth=2)
        self.canvas.draw()
        
    def remove_nearest_point(self, freq, vs):
        """删除离点击坐标最近的点"""
        if self.current_curve is None or len(self.current_curve) == 0:
            return
            
        # 找到最近的点
        distances = np.sqrt((self.current_curve[:, 0] - freq)**2)
        nearest_idx = np.argmin(distances)
        
        # 删除最近的点
        self.current_curve = np.delete(self.current_curve, nearest_idx, axis=0)
        
        # 更新存储
        if len(self.current_curve) == 0:
            self.current_curve = None
            if self.current_curve_index in self.curve_lines:
                try:
                    self.curve_lines[self.current_curve_index].remove()
                except (NotImplementedError, ValueError):
                    # 如果无法直接移除，将其设为不可见
                    self.curve_lines[self.current_curve_index].set_visible(False)
                # 清除标注
                for ann in self.point_annotations.get(self.current_curve_index, []):
                    try:
                        ann.remove()
                    except (NotImplementedError, ValueError):
                        # 如果无法直接移除，将其设为不可见
                        ann.set_visible(False)
                if self.current_curve_index in self.point_annotations:
                    del self.point_annotations[self.current_curve_index]
                del self.curve_lines[self.current_curve_index]
        else:
            self.curves[self.current_curve_index] = self.current_curve
            
        self.draw_curve()
        
    def switch_to_curve(self, curve_index):
        """切换到指定曲线"""
        self.current_curve_index = curve_index
        self.current_curve = self.curves.get(curve_index, np.empty((0, 2)))
        self.draw_curve()
        self.update_status()
        self.update_curve_buttons()  # 更新按钮高亮状态
        
    def new_curve(self):
        """创建新曲线"""
        new_index = max(self.curves.keys()) + 1 if self.curves else 1
        self.curves[new_index] = np.empty((0, 2))
        self.switch_to_curve(new_index)
        self.update_curve_buttons()
        
    def delete_current_curve(self):
        """删除当前曲线"""
        deleted_curve_index = self.current_curve_index
        if self.current_curve_index in self.curves:
            # 删除曲线的可视化元素
            if self.current_curve_index in self.curve_lines:
                try:
                    self.curve_lines[self.current_curve_index].remove()
                except (NotImplementedError, ValueError):
                    # 如果无法直接移除，将其设为不可见
                    self.curve_lines[self.current_curve_index].set_visible(False)
                # 清除标注
                for ann in self.point_annotations.get(self.current_curve_index, []):
                    try:
                        ann.remove()
                    except (NotImplementedError, ValueError):
                        # 如果无法直接移除，将其设为不可见
                        ann.set_visible(False)
                if self.current_curve_index in self.point_annotations:
                    del self.point_annotations[self.current_curve_index]
                del self.curve_lines[self.current_curve_index]
            
            # 删除曲线数据
            del self.curves[self.current_curve_index]
            self.current_curve = None
            
            # 切换到其他曲线（如果有的话）
            if self.curves:
                next_index = list(self.curves.keys())[0]
                self.switch_to_curve(next_index)
            else:
                # 如果没有其他曲线，创建一个新的
                self.curves[0] = np.empty((0, 2))
                self.switch_to_curve(0)
                
            self.update_curve_buttons()
        
    def save_curves(self):
        """保存拾取的曲线"""
        # 默认文件名
        default_file = ""
        if self.h5_file:
            # 基于H5文件名生成默认的disp文件名
            base_name = os.path.splitext(self.h5_file)[0]
            default_file = base_name + '.disp'
        
        # 让用户选择保存位置和文件名
        output_file, _ = QFileDialog.getSaveFileName(
            self, 'Save Curves', default_file, 'Disp files (*.disp)')
        
        # 如果用户取消了保存操作
        if not output_file:
            return
        
        try:
            with open(output_file, 'w') as f:
                f.write("# Picked dispersion curves\n")
                f.write("# Format: frequency velocity\n")
                for curve_idx, curve in self.curves.items():
                    if len(curve) > 0:
                        f.write(f"\n# Curve {curve_idx}\n")
                        for freq, vs in curve:
                            f.write(f"{freq:.6f} {vs:.6f}\n")
                            
            QMessageBox.information(self, "Success", f"Curves saved to {output_file}")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save curves: {str(e)}")
        
    def load_curves_from_file(self):
        """从文件加载曲线"""
        if self.h5_file:
            # 默认加载与h5文件同名的disp文件
            default_file = os.path.splitext(self.h5_file)[0] + '.disp'
        else:
            default_file = ''
            
        # 让用户选择文件
        disp_file, _ = QFileDialog.getOpenFileName(
            self, 'Load Curves', default_file, 'Disp files (*.disp)')
            
        if not disp_file or not os.path.exists(disp_file):
            return
            
        try:
            self.load_dispersion_curves(disp_file)
            QMessageBox.information(self, "Success", f"Curves loaded from {disp_file}")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to load curves: {str(e)}")
            
    def load_dispersion_curves(self, disp_file):
        """加载已保存的频散曲线"""
        curves = {}
        current_curve_index = None
        current_curve_points = []
        
        with open(disp_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith('# Curve'):
                    # 保存之前的曲线
                    if current_curve_index is not None and current_curve_points:
                        curves[current_curve_index] = np.array(current_curve_points)
                    
                    # 开始新的曲线
                    current_curve_index = int(line.split()[2])
                    current_curve_points = []
                elif line and not line.startswith('#'):
                    # 读取频率和速度数据
                    freq, vs = map(float, line.split())
                    # 将速度依附到self.f存在的频率点上
                    if hasattr(self, 'f') and len(self.f) > 0:
                        # 找到最接近的频率点
                        freq_idx = np.argmin(np.abs(self.f - freq))
                        freq = self.f[freq_idx]
                    current_curve_points.append([freq, vs])
            
            # 保存最后一条曲线
            if current_curve_index is not None and current_curve_points:
                curves[current_curve_index] = np.array(current_curve_points)
        
        # 更新曲线数据
        self.curves = curves
        if self.curves:
            # 设置当前曲线为第一条曲线
            first_curve_index = min(self.curves.keys())
            self.current_curve_index = first_curve_index
            self.current_curve = self.curves[first_curve_index]
        else:
            self.current_curve_index = 0
            self.current_curve = None
            
        # 绘制所有曲线
        for curve_index in sorted(self.curves.keys()):
            self.current_curve_index = curve_index
            self.current_curve = self.curves[curve_index]
            self.draw_curve()
        
        # 恢复到第一条曲线
        if self.curves:
            first_curve_index = min(self.curves.keys())
            self.switch_to_curve(first_curve_index)
            
        # 更新曲线按钮状态
        self.update_curve_buttons()
        self.update_status()
        
    def on_scroll(self, event):
        """处理滚轮事件，用于切换曲线"""
        # 只处理FV图上的滚轮事件
        if event.inaxes != self.ax_fv:
            return
            
        if len(self.curves) <= 1:
            return  # 只有一条曲线时不切换
            
        if event.button == 'up':
            # 向上滚动，切换到下一条曲线
            curve_indices = sorted(list(self.curves.keys()))
            current_index = curve_indices.index(self.current_curve_index)
            next_index = (current_index + 1) % len(curve_indices)
            self.switch_to_curve(curve_indices[next_index])
        elif event.button == 'down':
            # 向下滚动，切换到上一条曲线
            curve_indices = sorted(list(self.curves.keys()))
            current_index = curve_indices.index(self.current_curve_index)
            next_index = (current_index - 1) % len(curve_indices)
            self.switch_to_curve(curve_indices[next_index])

    def open_h5_file(self):
        """打开H5文件"""
        h5_file, _ = QFileDialog.getOpenFileName(
            self, 'Open H5 File', '', 'H5 files (*.h5 *.hdf5)')
            
        if not h5_file or not os.path.exists(h5_file):
            return
            
        # 保存当前曲线作为参考
        old_curves = {}
        for idx, curve in self.curves.items():
            old_curves[idx] = curve.copy()
            
        # 加载新文件
        self.h5_file = h5_file
        self.load_h5_data()
        self.plot_fv_and_waveforms()
        
        # 尝试加载对应的.disp文件
        self.try_load_disp_file()
        # 如果没有加载到曲线，则使用旧的曲线作为参考，在新的FV图上重新拾取
        if not self.curves and old_curves:
            self.update_curves_with_new_fv(old_curves)
        elif not self.curves:
            # 如果既没有加载到曲线，也没有旧曲线，则初始化参考曲线
            self.initialize_reference_curve()

    def show_fv_config_dialog(self):
        """显示FV配置对话框"""
        dialog = FVConfigDialog(
            self, 
            vs_min=self.vs[0], 
            vs_max=self.vs[-1], 
            f_min=self.fs, 
            f_max=self.fe
        )
        
        if dialog.exec() == QDialog.DialogCode.Accepted:
            values = dialog.get_values()
            self.fs = values['f_min']
            self.fe = values['f_max']
            # 重新生成速度数组
            self.vs = np.linspace(values['vs_min'], values['vs_max'], 800)
            # 如果有波形数据，重新计算FV
            if hasattr(self, 'traces') and len(self.traces) > 0:
                self.recalculate_fv()
                # 根据频率范围更新波形显示
                self.update_waveforms_by_frequency_range()
            # 重新绘制
            self.plot_fv_and_waveforms()
            # 重新绘制所有曲线
            for curve_index in sorted(self.curves.keys()):
                self.current_curve_index = curve_index
                self.current_curve = self.curves[curve_index]
                self.draw_curve()

    def recalculate_fv(self):
        """重新计算FV数据"""
        try:
            # 使用do_fv函数计算新的FV矩阵
            FV, f, vs = do_fv(
                self.traces, 
                self.vs, 
                x=self.x, 
                dt=self.dt, 
                fs=self.fs, 
                fe=self.fe, 
                NORM=self.CONFIG['FV_NORM']
            )
            # 更新FV数据
            self.FV = np.abs(FV)
            self.f = f
            self.vs = vs
            self.have_fv_data = True
        except Exception as e:
            print(f"Error recalculating FV: {e}")
            QMessageBox.critical(self, "Error", f"Failed to recalculate FV: {str(e)}")

    def update_waveforms_by_frequency_range(self):
        """根据频率范围更新波形显示"""
        if not (hasattr(self, 'traces') and hasattr(self, 't') and hasattr(self, 'x')):
            return
            
        # 计算需要滤波的频率范围
        f_min = max(self.fs, 0.1)  # 避免频率为0
        f_max = min(self.fe, 1.0 / (2 * self.dt))  # 不超过奈奎斯特频率
        
        # 创建临时副本用于显示
        self.filtered_traces = self.traces.copy()
        
        # 如果频率范围不是全频带，则对波形进行滤波处理
        if f_min > 0.1 or f_max < 1.0 / (2 * self.dt):
            try:
                # 使用简单的频域滤波
                nt = len(self.t)
                df = 1.0 / (self.t[-1] - self.t[0])
                frequencies = np.fft.fftfreq(nt, d=self.dt)
                
                # 创建频域滤波器
                filter_mask = np.zeros(nt, dtype=bool)
                filter_mask[(np.abs(frequencies) >= f_min) & (np.abs(frequencies) <= f_max)] = True
                
                # 对每个道进行滤波
                for i in range(len(self.filtered_traces)):
                    trace_fft = np.fft.fft(self.filtered_traces[i, :])
                    trace_fft[~filter_mask] = 0
                    self.filtered_traces[i, :] = np.real(np.fft.ifft(trace_fft))
                    
            except Exception as e:
                print(f"Warning: Failed to filter waveforms: {e}")
                # 如果滤波失败，使用原始数据
                self.filtered_traces = self.traces
                
        # 更新波形显示
        self.plot_waveforms()

    def try_load_disp_file(self):
        """尝试加载对应的.disp文件"""
        if not self.h5_file:
            return
            
        # 生成对应的.disp文件路径
        disp_file = os.path.splitext(self.h5_file)[0] + '.disp'
        
        # 检查文件是否存在
        if os.path.exists(disp_file):
            try:
                self.load_dispersion_curves(disp_file)
                print(f"Loaded dispersion curves from {disp_file}")
            except Exception as e:
                print(f"Error loading dispersion curves from {disp_file}: {e}")
        else:
            print(f"No corresponding .disp file found for {self.h5_file}")
            
    def update_curves_with_new_fv(self, old_curves):
        """基于新的FV图更新曲线"""
        # 清空当前曲线
        self.curves = {}
        
        # 对每条旧曲线，在新FV图上重新拾取
        for curve_index, old_curve in old_curves.items():
            new_curve_points = []
            for freq, old_vs in old_curve:
                # 在新FV图上，在旧速度附近搜索新的局部极大值
                new_vs = self.find_local_maximum(freq, old_vs)
                new_curve_points.append([freq, new_vs])
            
            # 保存更新后的曲线
            self.curves[curve_index] = np.array(new_curve_points)
        
        # 设置当前曲线
        if self.curves:
            self.current_curve_index = min(self.curves.keys())
            self.current_curve = self.curves[self.current_curve_index]
        else:
            self.current_curve_index = 0
            self.current_curve = None
            
        # 绘制所有曲线
        for curve_index in sorted(self.curves.keys()):
            self.current_curve_index = curve_index
            self.current_curve = self.curves[curve_index]
            self.draw_curve()
        
        # 恢复到第一条曲线
        if self.curves:
            first_curve_index = min(self.curves.keys())
            self.switch_to_curve(first_curve_index)
            
        # 更新曲线按钮状态
        self.update_curve_buttons()
        self.update_status()
        
    def save_figure(self):
        """保存当前图形为PNG文件"""
        # 默认文件名
        default_file = ""
        if self.h5_file:
            # 基于H5文件名生成默认的png文件名
            base_name = os.path.splitext(self.h5_file)[0]
            default_file = base_name + '.disp.png'
        
        # 让用户选择保存位置和文件名
        output_file, _ = QFileDialog.getSaveFileName(
            self, 'Save Figure', default_file, 'PNG files (*.png)')
        
        # 如果用户取消了保存操作
        if not output_file:
            return
            
        try:
            # 保存当前图形
            self.figure.savefig(output_file, dpi=self.CONFIG['DPI'], 
                               bbox_inches='tight', facecolor='white')
            QMessageBox.information(self, "Success", f"Figure saved to {output_file}")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save figure: {str(e)}")


def main():
    app = QApplication(sys.argv)
    
    # 获取命令行参数
    h5_file = None
    if len(sys.argv) > 1:
        h5_file = sys.argv[1]
    
    window = DispersionPickerWindow(h5_file)
    window.show()
    
    sys.exit(app.exec())


if __name__ == '__main__':
    main()